package middleware

import (
	"github.com/gin-gonic/gin"
	"io/ioutil"
	"oa_kadia_link/config"
	"oa_kadia_link/globe"
	"oa_kadia_link/utils"
	"time"
)

var tokenBucketHash string

func init() {
	globe.GetRdb().HMSet(utils.GetLimitBucketFlag(), map[string]interface{}{
		"current":                 config.GetCount().MaxLimitBucket,
		"maxLimitBucket":          config.GetCount().MaxLimitBucket,
		"onceLimitBucketIncr":     config.GetCount().OnceLimitBucketIncr,
		"limitBucketIncrInterval": config.GetCount().LimitBucketIncrInterval,
		"lastTime":                time.Now().Unix(),
	})
	tokenBucketScript, err := ioutil.ReadFile("./eval/tokenBucket.lua")
	utils.CheckError(err)
	tokenBucketHash, err = globe.GetRdb().ScriptLoad(string(tokenBucketScript)).Result()
	utils.CheckError(err)
}

func GetOneFromLimitBucket() bool {
	surplus, err := globe.GetRdb().EvalSha(tokenBucketHash, []string{utils.GetLimitBucketFlag()}, time.Now().Unix()).Result()
	utils.CheckError(err)
	return surplus.(int64) >= 0
}

func Limiting() gin.HandlerFunc {
	return func(context *gin.Context) {
		if !utils.IsAdmin(context) && !GetOneFromLimitBucket() {
			utils.ErrorJSON(context, "访问人数过多，请稍后再试！")
			context.Abort()
			return
		}
	}
}
